//
//  NetWorkSecurity.m
//  URLSessionTest
//
//  Created by mac on 2017/7/7.
//  Copyright © 2017年 yyb. All rights reserved.
//

#import "NetWorkSecurity.h"
#import "AssertMacros.h"

static id TWAFPublicKeyForCertificate(NSData *certificate) {
    id allowedPublicKey = nil;
    SecCertificateRef allowedCertificate;
    SecCertificateRef allowedCertificates[1];
    CFArrayRef tempCertificates = nil;
    SecPolicyRef policy = nil;
    SecTrustRef allowedTrust = nil;
    SecTrustResultType result;
    
    allowedCertificate = SecCertificateCreateWithData(NULL, (__bridge CFDataRef)certificate);
    __Require_Quiet(allowedCertificate != NULL, _out);
    
    allowedCertificates[0] = allowedCertificate;
    tempCertificates = CFArrayCreate(NULL, (const void **)allowedCertificates, 1, NULL);
    
    policy = SecPolicyCreateBasicX509();
    __Require_noErr_Quiet(SecTrustCreateWithCertificates(tempCertificates, policy, &allowedTrust), _out);
    __Require_noErr_Quiet(SecTrustEvaluate(allowedTrust, &result), _out);
    
    allowedPublicKey = (__bridge_transfer id)SecTrustCopyPublicKey(allowedTrust);
    
_out:
    if (allowedTrust) {
        CFRelease(allowedTrust);
    }
    
    if (policy) {
        CFRelease(policy);
    }
    
    if (tempCertificates) {
        CFRelease(tempCertificates);
    }
    
    if (allowedCertificate) {
        CFRelease(allowedCertificate);
    }
    
    return allowedPublicKey;
}

static BOOL TWAFServerTrustIsValid(SecTrustRef serverTrust) {
    BOOL isValid = NO;
    SecTrustResultType result;
    __Require_noErr_Quiet(SecTrustEvaluate(serverTrust, &result), _out);
    isValid = (result == kSecTrustResultUnspecified || result == kSecTrustResultProceed);
_out:
    return isValid;
}

static NSArray * TWAFCertificateTrustChainForServerTrust(SecTrustRef serverTrust) {
    CFIndex certificateCount = SecTrustGetCertificateCount(serverTrust);
    NSMutableArray *trustChain = [NSMutableArray arrayWithCapacity:(NSUInteger)certificateCount];
    
    for (CFIndex i = 0; i < certificateCount; i++) {
        SecCertificateRef certificate = SecTrustGetCertificateAtIndex(serverTrust, i);
        [trustChain addObject:(__bridge_transfer NSData *)SecCertificateCopyData(certificate)];
    }
    
    return [NSArray arrayWithArray:trustChain];
}


static NSArray * TWAFPublicKeyTrustChainForServerTrust(SecTrustRef serverTrust) {
    SecPolicyRef policy = SecPolicyCreateBasicX509();
    CFIndex certificateCount = SecTrustGetCertificateCount(serverTrust);
    NSMutableArray *trustChain = [NSMutableArray arrayWithCapacity:(NSUInteger)certificateCount];
    for (CFIndex i = 0; i < certificateCount; i++) {
        SecCertificateRef certificate = SecTrustGetCertificateAtIndex(serverTrust, i);
        
        SecCertificateRef someCertificates[] = {certificate};
        CFArrayRef certificates = CFArrayCreate(NULL, (const void **)someCertificates, 1, NULL);
        
        SecTrustRef trust;
        __Require_noErr_Quiet(SecTrustCreateWithCertificates(certificates, policy, &trust), _out);
        
        SecTrustResultType result;
        __Require_noErr_Quiet(SecTrustEvaluate(trust, &result), _out);
        
        [trustChain addObject:(__bridge_transfer id)SecTrustCopyPublicKey(trust)];
        
    _out:
        if (trust) {
            CFRelease(trust);
        }
        
        if (certificates) {
            CFRelease(certificates);
        }
        
        continue;
    }
    CFRelease(policy);
    
    return [NSArray arrayWithArray:trustChain];
}

static BOOL TWAFSecKeyIsEqualToKey(SecKeyRef key1, SecKeyRef key2) {
    return [(__bridge id)key1 isEqual:(__bridge id)key2];
}

@implementation NetWorkSecurity
- (instancetype)init {
    self = [super init];
    if (!self) {
        return nil;
    }
    self.validatesDomainName = YES;
    self.SSLPinningMode = TWAFSSLPinningModeNone;
    return self;
}

- (void)addSSLPinning:(NSData *)LocalCertData andAllowInvalidCertificates:(BOOL)allowInvalidCertificates andValidatesDomainName:(BOOL)validatesDomainName andSSLPinningMode:(TWAFSSLPinningMode )SSLPinningMode {
    self.allowInvalidCertificates = allowInvalidCertificates;
    self.validatesDomainName = validatesDomainName;
    self.SSLPinningMode = SSLPinningMode;
    self.pinnedCertificates  = [[NSSet alloc] initWithObjects:LocalCertData, nil];
}

- (void)setPinnedCertificates:(NSSet *)pinnedCertificates {
    _pinnedCertificates = pinnedCertificates;
    if (self.pinnedCertificates) {
        NSMutableSet *mutablePinnedPublicKeys = [NSMutableSet setWithCapacity:[self.pinnedCertificates count]];
        for (NSData *certificate in self.pinnedCertificates) {
            id publicKey = TWAFPublicKeyForCertificate(certificate);
            if (!publicKey) {
                continue;
            }
            [mutablePinnedPublicKeys addObject:publicKey];
        }
        self.pinnedPublicKeys = [NSSet setWithSet:mutablePinnedPublicKeys];
    } else {
        self.pinnedPublicKeys = nil;
    }
}

- (BOOL)evaluateServerTrust:(SecTrustRef)serverTrust
                  forDomain:(NSString *)domain {
    //当使用自建证书验证域名时，需要使用AFSSLPinningModePublicKey或者AFSSLPinningModeCertificate
    if (domain && self.allowInvalidCertificates && self.validatesDomainName && (self.SSLPinningMode == TWAFSSLPinningModeNone || [self.pinnedCertificates count] == 0)) {
        return NO;
    }
    //需要验证域名时，需要添加一个验证域名的策略
    NSMutableArray *policies = [NSMutableArray array];
    if (self.validatesDomainName) {
        [policies addObject:(__bridge_transfer id)SecPolicyCreateSSL(true, (__bridge CFStringRef)domain)];
    } else {
        [policies addObject:(__bridge_transfer id)SecPolicyCreateBasicX509()];
    }
    //设置验证的策略，可以是多个
    SecTrustSetPolicies(serverTrust, (__bridge CFArrayRef)policies);
    if (self.SSLPinningMode == TWAFSSLPinningModeNone) {
        return self.allowInvalidCertificates || TWAFServerTrustIsValid(serverTrust);
    } else if (!TWAFServerTrustIsValid(serverTrust) && !self.allowInvalidCertificates) {
        return NO;
    }
    //获取服务器证书的内容
    switch (self.SSLPinningMode) {
        case TWAFSSLPinningModeNone:
        default:
            return NO;
        case TWAFSSLPinningModeCertificate: {
            //直接将本地的证书设置为信任的根证书，然后来进行判断，并且比较本地证书的内容和服务器证书内容是否相同，如果有一个相同则返回YES
            NSMutableArray *pinnedCertificates = [NSMutableArray array];
            for (NSData *certificateData in self.pinnedCertificates) {
                [pinnedCertificates addObject:(__bridge_transfer id)SecCertificateCreateWithData(NULL, (__bridge CFDataRef)certificateData)];
            }
            //设置本地的证书为服务器的根证书
            SecTrustSetAnchorCertificates(serverTrust, (__bridge CFArrayRef)pinnedCertificates);
            //通过本地的证书来判断服务器证书是否可信，不可信，则验证不通过
            if (!TWAFServerTrustIsValid(serverTrust)) {
                NSLog(@"SSL证书错误,请检查本地证书是否正确");
                return NO;
            }
            // 获得正在验证的链条TWAFter，其中*应该包含最后位置的固定证书（如果是根CA）
            NSArray *serverCertificates = TWAFCertificateTrustChainForServerTrust(serverTrust);
            //客户端是否完全包含服务端证书
            for (NSData *trustChainCertificate in [serverCertificates reverseObjectEnumerator]) {
                if ([self.pinnedCertificates containsObject:trustChainCertificate]) {
                    return YES;
                }
            }
            return NO;
        }
        case TWAFSSLPinningModePublicKey: {
            //是通过比较证书当中公钥(PublicKey)部分来进行验证，通过SecTrustCopyPublicKey方法获取本地证书和服务器证书，然后进行比较，如果有一个相同，则通过验证
            NSUInteger trustedPublicKeyCount = 0;
            //服务端证书
            NSArray *publicKeys = TWAFPublicKeyTrustChainForServerTrust(serverTrust);
            //判断服务器证书的公钥与本地的证书公钥是否相同，相同则客户端认证通过
            for (id trustChainPublicKey in publicKeys) {
                //遍历客户端证书
                for (id pinnedPublicKey in self.pinnedPublicKeys) {
                    //对比两者的证书是否相同
                    if (TWAFSecKeyIsEqualToKey((__bridge SecKeyRef)trustChainPublicKey, (__bridge SecKeyRef)pinnedPublicKey)) {
                        trustedPublicKeyCount += 1;
                    }
                }
            }
            return trustedPublicKeyCount > 0;
        }
    }
    
    return NO;
}

@end
